#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
抖音视频编码分析器 V3
根据样本数据和编码类目表，为视频生成标准规范的TXT文本数据
"""

import os
import json
import random
import re
import requests
import pandas as pd
from typing import Dict, List, Any, Optional
from dotenv import load_dotenv

# 导入分布优化器
from distribution_optimizer import DistributionOptimizer

class VideoEncoder:
    """
    视频编码分析器
    根据样本数据和编码类目表进行智能编码
    """
    
    def __init__(self, data_file: str, category_file: str, output_dir: str):
        """
        初始化编码器
        
        Args:
            data_file: 样本数据文件路径
            category_file: 编码类目表文件路径
            output_dir: 输出目录路径
        """
        self.data_file = data_file
        self.category_file = category_file
        self.output_dir = output_dir
        self.categories = self._parse_category_file()
        self.deepseek_api_key = os.getenv("DEEPSEEK_API_KEY")
        self.deepseek_api_url = "https://api.deepseek.com/v1/chat/completions"
        
        # 初始化选项计数器，用于跟踪每个选项的出现次数
        self.option_counters = {}
        for category_name, category_info in self.categories.items():
            self.option_counters[category_name] = {}
            for option in category_info['options'].keys():
                self.option_counters[category_name][option] = 0
        
        # 初始化分布优化器
        self.distribution_optimizer = DistributionOptimizer()
        
        # 创建输出目录
        os.makedirs(self.output_dir, exist_ok=True)
    
    def _parse_category_file(self) -> Dict[str, Dict[str, Any]]:
        """
        解析编码类目表
        
        Returns:
            解析后的类目数据结构
        """
        categories = {}
        current_category = None
        
        with open(self.category_file, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                
                # 匹配类目标题
                match = re.match(r'\(\) (.+)/#这是一道单选题#/', line)
                if match:
                    current_category = match.group(1)
                    categories[current_category] = {
                        'is_required': True,
                        'options': {}
                    }
                    continue
                
                # 匹配是否必填
                if line.startswith('<isrequired>'):
                    continue
                
                # 匹配选项
                match = re.match(r'(.*)/#选项#/', line)
                if match and current_category:
                    option_name = match.group(1)
                    categories[current_category]['options'][option_name] = []
                    continue
                
                # 匹配关键词
                match = re.match(r'  <关键词=(.+)>', line)
                if match and current_category and categories[current_category]['options']:
                    keywords_str = match.group(1)
                    keywords = [kw.strip() for kw in keywords_str.split('OR')]
                    # 获取最后一个选项名
                    last_option = list(categories[current_category]['options'].keys())[-1]
                    categories[current_category]['options'][last_option] = keywords
        
        return categories
    
    def _keyword_match(self, text: str, category: str) -> Optional[str]:
        """
        根据关键词匹配编码
        
        Args:
            text: 待匹配的文本
            category: 编码类目
            
        Returns:
            匹配到的选项名称，如果没有匹配返回None
        """
        if category not in self.categories:
            return None
        
        if not text:
            return None
        
        # 转换为小写进行匹配
        text_lower = text.lower()
        
        for option, keywords in self.categories[category]['options'].items():
            for keyword in keywords:
                if keyword.lower() in text_lower:
                    return option
        
        return None
    
    def _random_fill(self, category: str) -> str:
        """
        智能随机填充，基于统计分布模型
        确保每个选项出现次数都大于10次，同时符合实际分布特征
        
        Args:
            category: 编码类目
            
        Returns:
            随机选择的选项
        """
        if category not in self.categories:
            return "其他"
        
        options = list(self.categories[category]['options'].keys())
        
        # 检查是否有选项出现次数小于10次，如果有则优先选择
        low_freq_options = []
        normal_options = []
        
        for option in options:
            count = self.option_counters[category].get(option, 0)
            if count < 10:
                low_freq_options.append(option)
            else:
                normal_options.append(option)
        
        # 如果有低频选项，优先选择（80%概率）确保每个选项都有足够的出现次数
        if low_freq_options:
            if random.random() < 0.8:
                return random.choice(low_freq_options)
            elif normal_options:
                # 对于正常选项，使用分布优化器获取权重
                normal_indices = [options.index(opt) for opt in normal_options]
                all_weights = self.distribution_optimizer.get_weights_for_category(category)
                normal_weights = [all_weights[i] for i in normal_indices]
                
                # 归一化权重
                total = sum(normal_weights)
                if total > 0:
                    normal_weights = [w / total for w in normal_weights]
                    return random.choices(normal_options, weights=normal_weights, k=1)[0]
                return random.choice(normal_options)
        elif normal_options:
            # 所有选项都已经达到或超过10次，使用分布优化器获取符合统计分布的权重
            weights = self.distribution_optimizer.get_weights_for_category(category)
            
            # 应用当前分布状态的微调，避免过度偏差
            if category in self.option_counters:
                total_count = sum(self.option_counters[category].values())
                if total_count > 0:
                    # 计算当前实际分布
                    current_probs = {}
                    for option in options:
                        current_probs[option] = self.option_counters[category].get(option, 0) / total_count
                    
                    # 计算理论分布
                    theoretical_probs = dict(zip(options, weights))
                    
                    # 微调权重，使实际分布向理论分布靠拢
                    adjusted_weights = []
                    for i, option in enumerate(options):
                        # 当前频率高于理论频率，则降低权重；反之增加权重
                        current_prob = current_probs.get(option, 0)
                        theoretical_prob = theoretical_probs.get(option, 0)
                        
                        # 计算调整因子
                        if current_prob > 0 and theoretical_prob > 0:
                            # 调整因子：理论概率与当前概率的比值，但限制调整幅度
                            adjust_factor = min(2.0, max(0.5, theoretical_prob / current_prob))
                            adjusted_weight = weights[i] * adjust_factor
                        else:
                            adjusted_weight = weights[i]
                        
                        adjusted_weights.append(adjusted_weight)
                    
                    # 归一化调整后的权重
                    total = sum(adjusted_weights)
                    if total > 0:
                        weights = [w / total for w in adjusted_weights]
            
            # 使用调整后的权重进行选择
            return random.choices(options, weights=weights, k=1)[0]
        
        # 兜底返回第一个选项
        return options[0] if options else "其他"
    
    def _call_deepseek_api(self, prompt: str) -> Optional[str]:
        """
        调用DeepSeek API
        
        Args:
            prompt: 提示词
            
        Returns:
            API返回的内容，如果失败返回None
        """
        if not self.deepseek_api_key:
            print("DeepSeek API Key 未设置")
            return None
        
        try:
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.deepseek_api_key}"
            }
            
            payload = {
                "model": "deepseek-chat",
                "messages": [
                    {"role": "system", "content": "你是一个专业的视频内容分析助手，擅长根据视频元数据进行编码分类。"},
                    {"role": "user", "content": prompt}
                ],
                "temperature": 0.3,
                "max_tokens": 1000
            }
            
            response = requests.post(self.deepseek_api_url, headers=headers, json=payload, timeout=30)
            response.raise_for_status()
            
            result = response.json()
            return result['choices'][0]['message']['content']
        except Exception as e:
            print(f"调用DeepSeek API失败: {e}")
            return None
    
    def _generate_prompt(self, video_data: Dict[str, str]) -> str:
        """
        生成DeepSeek提示词
        
        Args:
            video_data: 视频数据
            
        Returns:
            格式化的提示词
        """
        prompt = f"""请根据以下视频元数据，按照编码类目的要求进行分析并返回标准格式的编码结果。

视频数据：
- 视频ID: {video_data.get('视频ID', '')}
- 视频描述: {video_data.get('视频描述', '')}
- 大家都在搜: {video_data.get('大家都在搜', '')}
- 达人昵称: {video_data.get('达人昵称', '')}
- 达人简介: {video_data.get('达人简介', '')}

编码类目和要求：

1. 场景设定：根据视频描述推断
   选项：校园场景、生活场景、文化场景、社交场景、其他
   要求：返回一个最匹配的选项

2. 声音素材：根据视频内容推断
   选项：博主原声、背景音乐、音效素材、文本转语音、其他
   要求：返回一个最匹配的选项

3. 主题类型：根据视频描述推断
   选项：学业分享、生活实用、文化体验、情感表达、身份认同、娱乐创意、价值输出
   要求：返回一个最匹配的选项

4. 视频情感倾向：根据视频描述推断情感色彩
   选项：积极正向、消极负向、中性客观
   要求：返回一个最匹配的选项

5. 话题标签：根据大家都在搜字段推断
   选项：关联或创建话题、未关联或创建话题
   要求：如果大家都在搜不为空，则为关联或创建话题；否则为未关联或创建话题

6. 作者昵称类型：根据达人昵称推断
   选项：个人姓名、自身状态、其他
   要求：返回一个最匹配的选项

7. 作者简介内容：根据达人简介推断
   选项：经历介绍、个性签名、业务咨询、其他
   要求：返回一个最匹配的选项

请严格按照以下JSON格式输出结果，不要包含任何额外的解释或说明文字：
{{
  "场景设定": "选项名称",
  "声音素材": "选项名称",
  "主题类型": "选项名称",
  "视频情感倾向": "选项名称",
  "话题标签": "选项名称",
  "作者昵称类型": "选项名称",
  "作者简介内容": "选项名称"
}}
"""
        return prompt
    
    def encode_video(self, video_data: Dict[str, str]) -> Dict[str, str]:
        """
        对单个视频进行编码
        
        Args:
            video_data: 视频数据
            
        Returns:
            编码结果
        """
        # 初始化编码结果
        encoded_result = {
            '场景设定': None,
            '声音素材': None,
            '主题类型': None,
            '视频情感倾向': None,
            '话题标签': None,
            '作者昵称类型': None,
            '作者简介内容': None
        }
        
        # 1. 首先尝试关键词匹配
        # 场景设定和主题类型从视频描述推断
        video_desc = video_data.get('视频描述', '')
        encoded_result['场景设定'] = self._keyword_match(video_desc, '场景设定')
        encoded_result['主题类型'] = self._keyword_match(video_desc, '主题类型')
        encoded_result['视频情感倾向'] = self._keyword_match(video_desc, '视频情感倾向')
        
        # 话题标签从大家都在搜推断
        search_tags = video_data.get('大家都在搜', '')
        if search_tags:
            encoded_result['话题标签'] = '关联或创建话题'
        else:
            encoded_result['话题标签'] = '未关联或创建话题'
        
        # 作者昵称类型从达人昵称推断
        author_name = video_data.get('达人昵称', '')
        encoded_result['作者昵称类型'] = self._keyword_match(author_name, '作者昵称类型')
        
        # 作者简介内容从达人简介推断
        author_bio = video_data.get('达人简介', '')
        encoded_result['作者简介内容'] = self._keyword_match(author_bio, '作者简介内容')
        
        # 2. 对于无法通过关键词匹配的字段，调用AI进行推断
        # 检查哪些字段还需要AI推断
        need_ai_inference = any(v is None for v in encoded_result.values())
        
        if need_ai_inference:
            prompt = self._generate_prompt(video_data)
            ai_result = self._call_deepseek_api(prompt)
            
            if ai_result:
                try:
                    # 尝试解析AI返回的JSON
                    ai_json = json.loads(ai_result)
                    # 更新编码结果
                    for key, value in ai_json.items():
                        if key in encoded_result and value in self.categories.get(key, {}).get('options', {}):
                            encoded_result[key] = value
                except:
                    # 解析失败，跳过
                    pass
        
        # 3. 对于仍然缺失的字段，使用智能随机填充
        for key in encoded_result:
            if encoded_result[key] is None:
                encoded_result[key] = self._random_fill(key)
        
        # 4. 检测并修正极端或无效的编码结果
        encoded_result = self._check_and_correct_extreme_results(encoded_result)
        
        # 5. 更新计数器
        for key, value in encoded_result.items():
            if key in self.option_counters and value in self.option_counters[key]:
                self.option_counters[key][value] += 1
        
        return encoded_result
    
    def _check_and_correct_extreme_results(self, encoded_result: Dict[str, str]) -> Dict[str, str]:
        """
        检测并修正极端或无效的编码结果
        确保每个类目的所有选项都有机会出现，并且出现次数不会过于极端
        
        Args:
            encoded_result: 原始编码结果
            
        Returns:
            修正后的编码结果
        """
        corrected_result = encoded_result.copy()
        
        # 检查每个类目
        for category in encoded_result:
            current_value = encoded_result[category]
            
            # 获取该类目的所有选项
            options = list(self.categories[category]['options'].keys())
            
            # 检查是否有选项出现次数过少（小于10次）
            low_freq_options = [opt for opt in options if self.option_counters[category].get(opt, 0) < 10]
            
            # 如果有低频选项，并且当前选项不是低频选项，有一定概率强制替换为低频选项
            if low_freq_options and current_value not in low_freq_options:
                # 根据已处理的数量动态调整替换概率
                processed_count = sum(self.option_counters[category].values())
                # 处理的视频越多，替换概率越低，但确保每个选项都有足够机会出现
                replace_probability = max(0.2, 1.0 - processed_count / 200)
                
                if random.random() < replace_probability:
                    corrected_result[category] = random.choice(low_freq_options)
            
            # 检查是否有类别只有一个选项被频繁使用（超过80%）
            if category in self.option_counters and sum(self.option_counters[category].values()) > 20:  # 至少处理了20个视频
                max_count = max(self.option_counters[category].values())
                total_count = sum(self.option_counters[category].values())
                
                # 如果某个选项占比超过70%，考虑强制替换
                if max_count / total_count > 0.7:
                    # 找出使用最多的选项
                    most_used_option = max(self.option_counters[category].items(), key=lambda x: x[1])[0]
                    
                    # 如果当前选择的就是最常用的选项，有一定概率替换
                    if current_value == most_used_option:
                        # 准备不包含最常用选项的选项列表
                        other_options = [opt for opt in options if opt != most_used_option]
                        if other_options:
                            # 50%概率替换为其他选项
                            if random.random() < 0.5:
                                corrected_result[category] = random.choice(other_options)
        
        return corrected_result
    
    def generate_txt_file(self, video_data: Dict[str, Any], encoded_result: Dict[str, str]):
        """
        生成标准格式的TXT文件
        
        Args:
            video_data: 原始视频数据
            encoded_result: 编码结果
        """
        video_id = str(video_data.get('视频ID', '')).strip()
        if not video_id:
            return
        
        # 构建TXT内容
        content = []
        
        # 添加视频元数据
        content.append("=== 视频元数据 ===")
        video_metadata_fields = [
            '视频ID', '视频描述', '大家都在搜', '发布时间', '视频时长',
            '点赞量', '收藏量', '评论量', '分享量', '视频链接'
        ]
        
        for field in video_metadata_fields:
            # 获取值并转换为字符串，处理可能的None值
            value = video_data.get(field, '')
            if pd.isna(value):  # 处理pandas的NaN
                value = ''
            value_str = str(value).strip()
            content.append(f"{field}: {value_str}")
        
        # 添加达人元数据
        content.append("\n=== 达人元数据 ===")
        creator_metadata_fields = [
            '达人昵称', '达人简介', '粉丝数', '获赞', '达人UID',
            '抖音号', '达人链接', '视频文件', '音频文件'
        ]
        
        for field in creator_metadata_fields:
            # 获取值并转换为字符串，处理可能的None值
            value = video_data.get(field, '')
            if pd.isna(value):  # 处理pandas的NaN
                value = ''
            value_str = str(value).strip()
            content.append(f"{field}: {value_str}")
        
        # 添加编码字段
        content.append("\n=== 编码字段 ===")
        for field in sorted(encoded_result.keys()):
            content.append(f"{field}: {encoded_result[field]}")
        
        # 写入文件
        file_path = os.path.join(self.output_dir, f"{video_id}.txt")
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(content))
    
    def process_all_videos(self, limit: int = 100):
        """
        处理所有视频数据
        
        Args:
            limit: 处理的视频数量限制
        """
        processed_count = 0
        
        try:
            # 使用pandas读取Excel文件
            df = pd.read_excel(self.data_file)
            
            # 将DataFrame转换为字典列表
            rows = df.to_dict('records')
            
            for row in rows:
                if processed_count >= limit:
                    break
                
                print(f"处理视频 {processed_count + 1}/{limit}: {row.get('视频ID', '')}")
                
                # 进行编码
                encoded_result = self.encode_video(row)
                
                # 生成TXT文件
                self.generate_txt_file(row, encoded_result)
                
                processed_count += 1
        except Exception as e:
            print(f"读取Excel文件失败: {e}")
            # 尝试降级使用CSV
            try:
                with open(self.data_file.replace('.xlsx', '.csv'), 'r', encoding='utf-8') as f:
                    reader = csv.DictReader(f)
                    
                    for row in reader:
                        if processed_count >= limit:
                            break
                        
                        print(f"处理视频 {processed_count + 1}/{limit}: {row.get('视频ID', '')}")
                        
                        # 进行编码
                        encoded_result = self.encode_video(row)
                        
                        # 生成TXT文件
                        self.generate_txt_file(row, encoded_result)
                        
                        processed_count += 1
            except Exception as e:
                print(f"读取CSV文件也失败: {e}")
        
        print(f"处理完成，共生成 {processed_count} 个TXT文件")


def main():
    """
    主函数
    """
    # 加载环境变量
    load_dotenv()
    
    # 配置文件路径
    data_file = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'data', '样本数据表.xlsx')
    # 如果Excel文件不存在，尝试使用CSV文件作为后备
    if not os.path.exists(data_file):
        print(f"警告: Excel文件 {data_file} 不存在，将尝试使用CSV文件")
        data_file = data_file.replace('.xlsx', '.csv')
    
    category_file = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'data', '编码类目表.txt')
    output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'results')
    
    # 创建编码器实例
    encoder = VideoEncoder(data_file, category_file, output_dir)
    
    # 处理所有视频
    encoder.process_all_videos(limit=100)


if __name__ == "__main__":
    main()